#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import yaml
import os
import sys

SSH_PORT=22

# 定义组件允许的数量规则（方便后续增加）
ALLOWED_COMPONENT_COUNTS = {
    # 只能存在0或1个的组件
    "mysql": [0, 1],
    "nginx": [0, 1],
    "minio": [0, 1],
    "kibana": [0, 1],
    # 只能存在0、1或3个的组件
    "nacos": [0, 1, 3],
    "zookeeper": [0, 1, 3],
    "kafka": [0, 1, 3],
    "elasticsearch": [0, 1, 3],
    # 只能存在0、1或6个的组件
    "redis": [0, 1, 6],
    # 只能存在0、1、3、4、6个的组件
    "clickhouse": [0, 1, 3, 4, 6],
    # 只能存在0、1、3、4、5、6个的组件
    "hadoop": [0, 1, 3, 4, 5, 6]
}

def validate_component_count(component, count):
    """
    检查单个组件的数量是否有效
    
    参数:
        component: 组件名称
        count: 该组件的实例数量
    
    返回:
        如果数量有效返回True，无效返回False
    """
    if component not in ALLOWED_COMPONENT_COUNTS:
        print(f"警告: 组件 '{component}' 未定义在允许的规则中，跳过检查。")
        return True
    
    allowed_list = ALLOWED_COMPONENT_COUNTS[component]
    if count not in allowed_list:
        print(f"错误: 组件 '{component}' 的数量 {count} 无效。允许的数量: {allowed_list}")
        return False
    
    print(f"组件 '{component}' 的数量 {count} 有效。")
    return True

def validate_all_component_counts(component_hosts):
    """
    检查所有组件的数量是否有效
    
    参数:
        component_hosts: 字典，包含组件名到主机列表的映射
    
    说明:
        如果任何组件的数量无效，则退出程序
    """
    all_valid = True
    for component, hosts in component_hosts.items():
        count = len(hosts)
        if not validate_component_count(component, count):
            all_valid = False
    
    if not all_valid:
        print("错误: 一些组件的数量不符合要求，请检查配置。")
        sys.exit(1)
    else:
        print("所有组件数量检查通过。")

def main():
    # 检查命令行参数数量
    if len(sys.argv) < 3:
        print("Usage: python script.py <config_file> <output_file>")
        print("Example: python script.py config.yml deploy/inventory.ini")
        sys.exit(1)
    
    # 从命令行参数获取文件路径
    config_file = sys.argv[1]
    output_file = sys.argv[2]
    
    # 确保输出目录存在
    output_dir = os.path.dirname(output_file)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    try:
        # 读取并解析 YAML 文件
        with open(config_file, 'r') as stream:
            data = yaml.safe_load(stream)
    except FileNotFoundError:
        print(f"错误: 配置文件 {config_file} 不存在!")
        sys.exit(1)
    except yaml.YAMLError as exc:
        print(f"解析 YAML 文件时出错: {exc}")
        sys.exit(1)
    except Exception as exc:
        print(f"读取文件时发生未知错误: {exc}")
        sys.exit(1)
    
    # 检查hosts字段
    if 'hosts' not in data or not isinstance(data['hosts'], list) or len(data['hosts']) == 0:
        print("错误: 配置文件中未找到有效的 'hosts' 列表!")
        sys.exit(1)
    
    # 解析所有主机信息
    host_list = []
    for host_str in data['hosts']:
        if not isinstance(host_str, str):
            print(f"警告: 跳过非字符串主机信息: {host_str}")
            continue
        
        parts = host_str.split(';')
        if len(parts) < 5:  # 现在需要至少5个部分
            print(f"警告: 主机信息格式不正确，需要至少5个部分: {host_str}")
            continue
        
        ip = parts[0].strip()
        hostname = parts[1].strip()  # 新增的hostname字段
        username = parts[2].strip()
        password = parts[3].strip()
        services_str = parts[4].strip()
        
        # 处理中文逗号和不规范分隔符
        services_str = services_str.replace('，', ',').replace('，', ',')
        services = [s.strip() for s in services_str.split(',') if s.strip()]
        
        host_list.append({
            'ip': ip,
            'hostname': hostname,  # 添加hostname到主机信息
            'username': username,
            'password': password,
            'services': services
        })
    
    if not host_list:
        print("错误: 未解析到任何有效主机信息!")
        sys.exit(1)
    
    # 构建组件到主机的映射
    component_hosts = {}
    for host in host_list:
        for service in host['services']:
            if service not in component_hosts:
                component_hosts[service] = []
            component_hosts[service].append(host)

    # 添加组件数量验证
    validate_all_component_counts(component_hosts)
    
    # 定义多节点组件列表
    multi_node_components = ["nacos", "hadoop", "zookeeper", "kafka", "elasticsearch", "clickhouse", "redis"]

    # 定义组件的默认端口
    component_ports = {
        "nginx": 8080,
        "redis": 6379,
        "minio": 40169,
        "mysql": 3306,
        "zookeeper": 2181,
        "kafka": 9092,
        "elasticsearch": 9200,
        "kibana": 5601,
        "clickhouse": 8123,
        "nacos": 8848
        # hadoop 没有默认端口，因此不包含在此字典中
    }
    
    # 定义组件及其额外变量模板
    component_extra_vars = {
        "nginx": "nginx_mode={mode}",
        "redis": "redis_mode={mode} port=6379",
        "minio": "minio_mode={mode}",
        "mysql": "mysql_mode={mode}",
        "zookeeper": "zookeeper_mode={mode} zk_id={index}",
        "kafka": "kafka_mode={mode} kafka_id={index}",
        "elasticsearch": "es_mode={mode} master=true",
        "kibana": "kibana_mode={mode}",
        "clickhouse": "ck_mode={mode} shard_id={shard_id} keeper_id={keeper_id}",
        "nacos": "nacos_mode={mode}",
        "hadoop": "hadoop_mode={mode} hostname={hostname}"
    }
    
    # 生成 INI 文件内容
    try:
        with open(output_file, 'w') as f:
            for component, extra_var_template in component_extra_vars.items():
                if component not in component_hosts:
                    continue
                
                hosts_for_component = component_hosts[component]
                num_hosts = len(hosts_for_component)
                mode = "single" if num_hosts == 1 else "cluster"
                
                # 计算 addr_value：如果组件有默认端口，则构建所有节点的 IP:端口字符串
                port = component_ports.get(component)
                if port is not None:
                    addr_value = ",".join([f"{host['ip']}:{port}" for host in hosts_for_component])
                else:
                    addr_value = None  # 对于没有端口的组件（如 hadoop），不添加 addr

                f.write(f"[{component}]\n")
                
                # 特殊处理：Hadoop组件
                if component == "hadoop":
                    for index, host in enumerate(hosts_for_component):
                        # 基础变量，所有节点都有
                        base_vars = f"hadoop_mode={mode} hostname={host['hostname']} jdk=true"
                        
                        # 根据节点索引动态添加特定变量
                        if index < 2: # 前2个节点
                            extra_vars = f"{base_vars} jn=true rpc_port=8020 http_port=9870"
                        elif index == 2: # 第3个节点
                            extra_vars = f"{base_vars} jn=true"
                        else: # 第4个及以后的节点
                            extra_vars = base_vars
                        
                        # 写入这一行 (Hadoop组件的写入在这里完成)
                        f.write(f"{component}_node{index+1} ansible_host={host['ip']} ansible_port={SSH_PORT} ansible_user={host['username']} ansible_password='{host['password']}' {extra_vars}\n")
                
                # 其他组件的处理
                else:
                    for index, host in enumerate(hosts_for_component):
                        # 初始化 shard_id 和 keeper_id 的默认值
                        shard_id_val = index + 1
                        keeper_id_val = index + 1
                        
                        # 特殊处理：ClickHouse 组件
                        if component == "clickhouse" and mode == "cluster":
                            shard_id_val = (index // 2) + 1
                            if index < 3:
                                keeper_id_val = index + 1
                            else:
                                keeper_id_val = ""
                        
                        # 格式化额外变量
                        if component == "clickhouse":
                            extra_vars = extra_var_template.format(
                                mode=mode, 
                                shard_id=shard_id_val, 
                                keeper_id=keeper_id_val
                            )
                            if keeper_id_val == "":
                                extra_vars = extra_vars.replace("keeper_id=", "")
                        else:
                            extra_vars = extra_var_template.format(mode=mode, index=index)
                        
                        # 如果 addr_value 存在（即组件有端口），则添加到 extra_vars
                        if addr_value is not None:
                            extra_vars += f" addr={addr_value}"

                        # 写入其他组件的每一行
                        f.write(f"{component}_node{index+1} ansible_host={host['ip']} ansible_port={SSH_PORT} ansible_user={host['username']} ansible_password='{host['password']}' {extra_vars}\n")
                
                f.write("\n") # 组件间空行
            
        print(f"成功: INI 文件已生成: {output_file}")
    except IOError as e:
        print(f"错误: 写入文件 {output_file} 时出错: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()
